﻿//#define VERBOSE

using System;
using System.Diagnostics;
using RayDen.Library.Components.Surface;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.Library.Entity.Scene;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Engines.PathTracer;
using RayDen.RayEngine.Samplers;

namespace RayDen.RayEngine.Engines.MCMC
{
    // ReSharper disable RedundantThisQualifier
    
    public class MltPathTracerPath : PathSamplerBase
    {
        internal RgbSpectrum Throughput;
        public PathTracerPathState PathState;


        protected int depth, tracedShadowRayCount;
        protected float pathWeight;
        protected float lastPdfW;
        protected bool specularBounce;

        public UInt64 PathExpression;

        protected RayEngineScene scene;

        public int MaxRaysPerPath
        {
            get;
            set;
        }

        public class PathVertex
        {
            public Point Point;
            public Vector OutgoingDirection, IncomingDirection;
            public float VertexWeight;

            public int HitItemId;
        }

        public class EyePathVertex : PathVertex
        {
            public float ImageX, ImageY, LensU, LensV;
        }

        public class ImplicitLightVertex : PathVertex
        {
            public ILight Light;
            public RgbSpectrum Sample;
        }

        public class EnvironmentPathVertex : PathVertex
        {

        }

        public class ExplicitLightVertex : PathVertex
        {
            public ILight Light;
            public RgbSpectrum Sample;
            public float SampleWeight;
            public float u0, u1, u2, u3;
        }

        public class GeometryVertex : PathVertex
        {
            public Normal GeometryNormal, ShadingNormal;
            public ISurfaceMaterial SampledBsdf;
            public RgbSpectrum Sample;
            public float BsdfSampleWeight;
            public float u0, u1, u2, u3;

            public ExplicitLightVertex LightSample;
        }

        public class PathVertexFactory
        {

            protected RayEngineScene scene;
            protected SurfaceSampler surfaceSampler;


            public PathVertex CreatePathVertex(ref RayData ray, ref RayHit rayHit, Sample Sample)
            {
                PathVertex result = null;
                bool missed = rayHit.Index == 0xffffffffu;
                if (missed)
                {
                    if (scene.EnvironmentMap == null)
                    {
                        result = new EnvironmentPathVertex() { IncomingDirection = -ray.Dir };
                    }
                    else
                    {
                        result = new ImplicitLightVertex() { IncomingDirection = -ray.Dir, Light = scene.EnvironmentMap };
                    }
                    return result;
                }
                var hitInfo = surfaceSampler.GetIntersection(ref ray, ref rayHit);
                var hitPoint = ray.Point(rayHit.Distance);

                var currentTriangleIndex = (int)rayHit.Index;
                if (hitInfo.IsLight)
                {
                    var lt = scene.GetLightByIndex(currentTriangleIndex);
                    result = new ImplicitLightVertex()
                        {
                            IncomingDirection = -ray.Dir,
                            Light = lt,
                            HitItemId = currentTriangleIndex,
                            Point = hitPoint
                        };
                }
                else
                {
                    Vector wo = -ray.Dir;
                    float fPdf;
                    Vector wi;
                    bool specularBounce;
                    float u0 = Sample.GetLazyValue(), u1 = Sample.GetLazyValue(), u2 = Sample.GetLazyValue(), u3 = Sample.GetLazyValue();


                    var bsdf = hitInfo.MMaterial;
                    var throughput = RgbSpectrum.Unit;

                    RgbSpectrum f = bsdf.Sample_f(ref wo, out wi, ref hitInfo.Normal,
                                                               ref hitInfo.ShadingNormal, ref throughput,
                                                               u0, u1, u2, ref hitInfo.TextureData,
                                                               out fPdf, out specularBounce);
                    result = new GeometryVertex()
                        {
                            BsdfSampleWeight = fPdf,
                            GeometryNormal = hitInfo.Normal,
                            ShadingNormal = hitInfo.ShadingNormal,
                            HitItemId = currentTriangleIndex,
                            IncomingDirection = wo,
                            OutgoingDirection = wi,
                            Point = hitPoint,
                            Sample = f,
                            u0 = u0,
                            u1 = u1,
                            u2 = u2,
                            u3 = u3,
                            SampledBsdf = bsdf
                        };

                }
                return result;
            }


            public EyePathVertex CreateEyePathVertex(Sample sample)
            {
                IRay ray;
                scene.Camera.GetRay(sample.imageX, sample.imageY, out ray);
                return new EyePathVertex()
                    {
                        ImageX = (float)sample.imageX,
                        ImageY = (float)sample.imageY,
                        Point = ray.Origin,
                        OutgoingDirection = ray.Direction,
                        VertexWeight = 1f,
                    };
            }

            public ExplicitLightVertex CreateLightVertex(Sample sample, GeometryVertex vertex)
            {

                int currentLightIndex = scene.SampleLights(sample.GetLazyValue());
                var light = scene.Lights[currentLightIndex];

                var ls = new LightSample();
                light.EvaluateShadow(ref vertex.Point, ref vertex.GeometryNormal, sample.GetLazyValue(), sample.GetLazyValue(), sample.GetLazyValue(), ref ls);
                if (ls.Pdf <= 0f)
                {
                    throw new Exception();
                }
                /*
                secRays[tracedShadowRayCount].color = new RgbSpectrum(ls.Spectrum);
                secRays[tracedShadowRayCount].pdf = ls.Pdf;
                secRays[tracedShadowRayCount].shadowRay = ls.LightRay;

                Vector lwi = secRays[tracedShadowRayCount].shadowRay.Dir;
                RgbSpectrum fs;
                hitInfo.MMaterial.f(
                    ref secRays[tracedShadowRayCount].shadowRay.Dir,
                    ref wo, ref hitInfo.ShadingNormal, out fs, BrdfType.Diffuse);
                secRays[tracedShadowRayCount].color *= lightTroughtput *
                                                       Vector.AbsDot(ref hitInfo.ShadingNormal, ref lwi) *
                                                       fs;
                */
                return new ExplicitLightVertex()
                    {
                        
                    };
            }
        }

        public class LightPathInfo
        {
            public int ID;
            public UInt64 PathExpressionData;
            public PathVertex[] Vertices;
        }


        internal struct ShadowRayInfo
        {
            public float pdf;
            public RgbSpectrum color;
            public RayData shadowRay;
            public int currentShadowRayIndex;
        }
        private ShadowRayInfo[] secRays;

        float Mis(float aPdf)
        {
            return aPdf;
        }

        // Mis weight for 2 pdfs
        float Mis2(
            float aSamplePdf,
            float aOtherPdf)
        {
            return Mis(aSamplePdf) / (Mis(aSamplePdf) + Mis(aOtherPdf));
        }

        public override void InitPath(IPathProcessor buffer)
        {
            base.InitPath(buffer);
            this.scene = pathIntegrator.Scene;
            this.Radiance = new RgbSpectrum(0f);
            this.Throughput = new RgbSpectrum(1f);
            this.PathState = PathTracerPathState.EyeVertex;
            this.secRays = new ShadowRayInfo[scene.ShadowRayCount];
            this.Sample = pathIntegrator.Sampler.GetSample(null);
            IRay ray;
            pathIntegrator.Scene.Camera.GetRay(Sample.imageX, Sample.imageY, out ray);
            this.PathRay = (RayData)ray;
            this.RayIndex = -1;
            this.pathWeight = 1.0f;
            this.lastPdfW = 1.0f;
            this.tracedShadowRayCount = 0;
            this.depth = 0;
            this.specularBounce = true;
        }

        public override bool FillRayBuffer(RayBuffer rayBuffer)
        {
            var leftSpace = rayBuffer.LeftSpace();
            if (((PathState == PathTracerPathState.EyeVertex) && (1 > leftSpace)) ||
            ((PathState == PathTracerPathState.ShadowRaysOnly) && (tracedShadowRayCount > leftSpace)) ||
            ((PathState == PathTracerPathState.NextVertex) && (tracedShadowRayCount + 1 > leftSpace)))
                return false;
            if (PathState != PathTracerPathState.ShadowRaysOnly)
                RayIndex = rayBuffer.AddRay(PathRay);
            if (PathState == PathTracerPathState.NextVertex || PathState == PathTracerPathState.ShadowRaysOnly)
            {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                    secRays[i].currentShadowRayIndex = rayBuffer.AddRay(secRays[i].shadowRay);
            }
            return true;
        }

        public override void Splat(SampleBuffer sampleBuffer)
        {
            sampleBuffer.SplatSample(this.Sample.imageX, this.Sample.imageY, ref this.Radiance);
        }

        public override void Advance(RayBuffer rayBuffer, SampleBuffer consumer)
        {
            int currentTriangleIndex = 0;
            SurfaceIntersectionData hitInfo = null;
#if VERBOSE

            try {
#endif
            var misWeight = 1.0f;

            base.Advance(rayBuffer, consumer);

            RayHit rayHit = rayBuffer.rayHits[RayIndex];

            if (((PathState == PathTracerPathState.ShadowRaysOnly) || (PathState == PathTracerPathState.NextVertex)) &&
                (tracedShadowRayCount > 0))
            {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                {
                    RayHit shadowRayHit = rayBuffer.rayHits[secRays[i].currentShadowRayIndex];
                    RgbSpectrum attenuation;
                    if (this.ShadowRayTest(ref shadowRayHit, out attenuation))
                    {
                        //                            Radiance.MADD()
                        Radiance += pathWeight * attenuation * ((secRays[i].color) / secRays[i].pdf);
                        pathWeight *= secRays[i].pdf;
                    }
                }
                Splat(consumer);
                
                return;
            }

            depth++;
            bool missed = rayHit.Index == 0xffffffffu;
            if (missed || PathState == PathTracerPathState.ShadowRaysOnly || depth > scene.MaxPathDepth)
            {
                if (depth > 0 && !specularBounce)
                {
                    misWeight = Mis2(lastPdfW, 1f / MathLab.INVPI);
                }

                if (specularBounce && missed)
                {
                    Radiance += this.SampleEnvironment(PathRay.Dir) * Throughput * pathWeight * misWeight;
                }

                Splat(consumer);
                
                return;
            }

            // Something was hit
            hitInfo = SurfaceSampler.GetIntersection(ref PathRay, ref rayHit);
            currentTriangleIndex = (int)rayHit.Index;

            if (hitInfo == null)
            {
                Debugger.Break();
            }
            //If Hit light)
            if (hitInfo.IsLight)
            {
                float pw = 1.0f;
                if (!specularBounce)
                {
                    pw *= pathWeight;
                }
                {
                    var lt = scene.GetLightByIndex(currentTriangleIndex);
                    if (lt != null)
                    {
                        var le = (RgbSpectrum)(lt.Le(ref PathRay.Dir));
                        Radiance += Throughput * le * pw;
                    }
                }
                Splat(consumer);
                
                return;
            }

            var hitPoint = PathRay.Point(rayHit.Distance);

            Vector wo = -PathRay.Dir;
            tracedShadowRayCount = 0;
            if (hitInfo.MMaterial.IsDiffuse())
            {
                float lightStrategyPdf = scene.ShadowRayCount / (float)scene.Lights.Length;
                RgbSpectrum lightTroughtput = Throughput * hitInfo.Color;
                for (int i = 0; i < scene.ShadowRayCount; ++i)
                {
                    int currentLightIndex = scene.SampleLights(Sample.GetLazyValue());
                    var light = scene.Lights[currentLightIndex];

                    var ls = new LightSample();
                    light.EvaluateShadow(ref hitPoint, ref hitInfo.ShadingNormal, Sample.GetLazyValue(), Sample.GetLazyValue(), Sample.GetLazyValue(), ref ls);
                    if (ls.Pdf <= 0f)
                    {
                        continue;
                    }

                    secRays[tracedShadowRayCount].color = (RgbSpectrum)(ls.Spectrum);
                    secRays[tracedShadowRayCount].pdf = ls.Pdf;
                    secRays[tracedShadowRayCount].shadowRay = ls.LightRay;

                    Vector lwi = secRays[tracedShadowRayCount].shadowRay.Dir;
                    RgbSpectrum fs;
                    hitInfo.MMaterial.f(
                        ref secRays[tracedShadowRayCount].shadowRay.Dir,
                        ref wo, ref hitInfo.ShadingNormal, ref Throughput, out fs, types: BrdfType.Diffuse);
                    secRays[tracedShadowRayCount].color *= lightTroughtput *
                                                           Vector.AbsDot(ref hitInfo.ShadingNormal, ref lwi) *
                                                           fs;
                    if (!secRays[tracedShadowRayCount].color.IsBlack())
                    {
                        secRays[tracedShadowRayCount].pdf *= lightStrategyPdf;
                        tracedShadowRayCount++;
                    }
                }
            }

            float fPdf;
            Vector wi;

            RgbSpectrum f = hitInfo.MMaterial.Sample_f(ref wo, out wi, ref hitInfo.Normal, ref hitInfo.ShadingNormal, ref Throughput,
                                                      Sample.GetLazyValue(), Sample.GetLazyValue(),
                                                      Sample.GetLazyValue(), ref hitInfo.TextureData,
                                                      out fPdf, out specularBounce);

            if ((fPdf <= 0.0f) || f.IsBlack())
            {
                if (tracedShadowRayCount > 0)
                    PathState = PathTracerPathState.ShadowRaysOnly;
                else
                {
                    Splat(consumer);
                    
                }
                return;
            }
            pathWeight *= fPdf;
            Throughput *= f / fPdf;

            if (depth > scene.MaxPathDepth)
            {
                float prob = Math.Max(Throughput.Filter(), scene.RussianRuletteImportanceCap);
                if (prob >= Sample.GetLazyValue())
                {
                    Throughput /= prob;
                    pathWeight *= prob;
                }
                else
                {
                    if (tracedShadowRayCount > 0)
                        PathState = PathTracerPathState.ShadowRaysOnly;
                    else
                    {
                        Splat(consumer);
                        
                    }

                    return;
                }
            }

            PathRay.Org = hitPoint;
            PathRay.Dir = wi.Normalize();
            PathState = PathTracerPathState.NextVertex;

#if VERBOSE
            }
            catch (Exception ex) {
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Advance path exception");
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Error triangle {0}", currentTriangleIndex);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.Message);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.StackTrace);

            }
#endif
        }


        private bool ShadowRayTest(ref RayHit shadowRayHit, out RgbSpectrum attenuation)
        {
            var hit = shadowRayHit.Index == 0xffffffffu || scene.IsLight((int)shadowRayHit.Index);
            attenuation = new RgbSpectrum(1f);
            if (hit) return true;
            var mesh = scene.GetMeshByTriangleIndex((int)shadowRayHit.Index);
            var mat =
                //SurfaceMaterials.CreateMaterial(scene.MaterialProvider.Get(mesh.MaterialName));
                scene.MatLib.GetSurfMat(mesh != null ? mesh.MaterialName : "default");

            if (mat.AlphaTexture != null)
            {
                attenuation = (RgbSpectrum)(mat.AlphaTexture as ConstTextureInfo).Color;
                hit = true;
                //mat.MaterialData.Kt.Filter() < Sample.GetLazyValue();
            }

            return hit;
        }

        private RgbSpectrum SampleEnvironment(Vector vector)
        {
            return this.scene.SampleEnvironment(-vector);
        }
    }
    // ReSharper restore RedundantThisQualifier
}